package com.xs.langchain4j_springboot.test;

import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.LocalDate;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;

import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;
import static java.util.concurrent.TimeUnit.SECONDS;
@RestController
public class ServiceWithStreamingExample {

    interface Assistant {
        Flux<String> chat(String message);
        TokenStream  chat2(String message);
    }


    // Sorry, "demo" API key does not support streaming (yet). Please use your own key.
    OpenAiStreamingChatModel model = OpenAiStreamingChatModel.builder()
            .apiKey("26f9ab1f-3509-4f8e-b06e-7495bc93373a")
            .modelName("doubao-1.5-pro-32k-250115")
            .baseUrl("https://ark.cn-beijing.volces.com/api/v3")
            .build();

    Assistant assistant = AiServices.create(Assistant.class, model);

    Flux<String> tellMeAJoke = assistant.chat("Tell me a joke");


    @RequestMapping(value = "/streammm",produces ="text/stream;charset=UTF-8")
    public Flux<String> memoryStreamChat(@RequestParam String message) {
        return assistant.chat(message);
    }

    @RequestMapping(value = "/streammm2",produces ="text/stream;charset=UTF-8")
    public void memoryStreamChat2() throws ExecutionException, InterruptedException, TimeoutException {

        OpenAiStreamingChatModel model = OpenAiStreamingChatModel.builder()
                .apiKey("26f9ab1f-3509-4f8e-b06e-7495bc93373a")
                .modelName("doubao-1.5-pro-32k-250115")
                .baseUrl("https://ark.cn-beijing.volces.com/api/v3")
                .build();
        Assistant assistant = AiServices.create(Assistant.class, model);

        TokenStream tokenStream = assistant.chat2("Tell me a joke");

        CompletableFuture<ChatResponse> futureResponse = new CompletableFuture<>();

        tokenStream.onPartialResponse(System.out::print)
                .onCompleteResponse(futureResponse::complete)
                .onError(futureResponse::completeExceptionally)
                .start();

        ChatResponse chatResponse = futureResponse.get(30, SECONDS);
        System.out.println("\n" + chatResponse);

    }
}